"""
RFdiffusion Hydra 回调类
"""

import os
import logging
from omegaconf import DictConfig
from hydra.core.callbacks import Callback

class SetBackendCallback(Callback):
    """设置图后端环境变量的回调类"""
    
    def __init__(self, backend_name="dgl", verbose=True):
        self.backend_name = backend_name
        self.verbose = verbose
        self.log = logging.getLogger(self.__class__.__name__)
    
    def on_job_start(self, config: DictConfig, **kwargs) -> None:
        """在作业开始时设置环境变量"""
        # 验证后端名称
        valid_backends = ["dgl", "pyg"]
        backend = self.backend_name.lower()
        
        if backend not in valid_backends:
            self.log.warning(f"无效的图后端 '{backend}'，使用默认值 'dgl'")
            backend = "dgl"
        
        # 设置统一的环境变量
        os.environ["GRAPH_BACKEND"] = backend
        
        # 同时设置两个组件的后端环境变量，确保兼容性
        os.environ["SE3_TRANSFORMER_GRAPH_BACKEND"] = backend
        os.environ["RFDIFFUSION_GRAPH_BACKEND"] = backend
        
        if self.verbose:
            self.log.info(f"设置图后端为: {backend}")
            
        return None 